[v1][core] Support for attention free models #20811

christian-pinto · 2025-07-11T10:55:37Z

This PR enables support for attention free models by exploiting the the zero kv cache groups approach. This han been created as a follow-up from the discussions in #20577.

@heheda12345 please review.

gemini-code-assist

Summary of Changes

Hello @christian-pinto, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces comprehensive support for 'attention-free models' by adapting the KV cache management system. My changes ensure that components dealing with KV cache allocation, configuration, and memory checks correctly handle models that do not require a KV cache, preventing errors and optimizing resource management for these specific model types. This is achieved by explicitly checking for and handling cases where KV cache groups are zero or empty.

Highlights

Core KV Cache Management Adaptation: I've updated the KVCacheCoordinator and KVCacheManager to gracefully handle models that do not utilize a KV cache (attention-free models). This involves adding checks to ensure that operations dependent on the existence of KV cache groups are only performed when such groups are present, preventing errors when kv_cache_groups is empty.
KV Cache Utility Enhancements: I've introduced new utility functions in kv_cache_utils.py to specifically identify and configure attention-free models. This includes is_kv_cache_type_attention_free to detect such models and _get_kv_cache_config_attention_free to provide a minimal KVCacheConfig suitable for them, which is then prioritized in get_kv_cache_config.
Memory and Spec Determination Optimization: I've optimized the AbstractWorker in abstract.py to skip unnecessary memory determination and KV cache specification retrieval for attention-free models. Since these models don't use a KV cache, these operations are redundant and can be bypassed, returning default or zero values.
Memory Check Bypass: I've added a condition to check_enough_kv_cache_memory in kv_cache_utils.py to bypass the memory availability check entirely if the model is identified as attention-free. This is because attention-free models do not require memory allocation for KV cache blocks.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces support for attention-free models by handling cases where the KV cache configuration is empty. I've identified a correctness issue in a new utility function and a couple of typos in comments. Once these are addressed, the PR should be in good shape.

gemini-code-assist · 2025-07-11T10:56:52Z

vllm/v1/core/kv_cache_utils.py

+    if not kv_cache_spec:
+        return True


This function is type-hinted to return a bool, but it implicitly returns None if kv_cache_spec is not empty. This violates the function's contract and could lead to unexpected behavior. Replace the if statement with a direct return to fix this.

return not kv_cache_spec

@christian-pinto please fix it

gemini-code-assist · 2025-07-11T10:56:52Z

vllm/v1/core/kv_cache_utils.py

+        # This returns a kv_cahce config with 0 kv cache groups and 1 block
+        # to allow for the KVCache manager to handle attention fre models.


There are a couple of typos in this comment: kv_cahce should be kv_cache, and attention fre should be attention free.

# This returns a kv_cache config with 0 kv cache groups and 1 block # to allow for the KVCache manager to handle attention-free models.

christian-pinto · 2025-07-11T11:00:49Z

vllm/v1/core/kv_cache_coordinator.py

-        self.verify_and_split_kv_cache_groups()
+        # attention free models are initialized with 0 kv_cache_groups
+        if len(self.kv_cache_config.kv_cache_groups) > 0:
+            self.verify_and_split_kv_cache_groups()


I will revert this once #20661 is merged and the KVCacheCoordinatorNoPrefixCache is available.

github-actions · 2025-07-11T11:05:00Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

heheda12345

Thank you for the great work. Let's wait for #20661.

heheda12345 · 2025-07-11T12:08:53Z

vllm/v1/core/kv_cache_utils.py

+    if not kv_cache_spec:
+        return True


@christian-pinto please fix it

heheda12345 · 2025-07-11T12:13:44Z

vllm/v1/core/kv_cache_manager.py

@@ -86,7 +86,7 @@ def __init__(
        self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

        self.block_size: Optional[int] = None
-        if self.enable_caching:
+        if self.enable_caching and len(self.kv_cache_config.kv_cache_groups) > 0:


Can you fix the type error?

All looks good to me here. Am i missing something?

I can see some precommit error

Oh I see, I have missed this one.
I'll take care of it. Thanks

heheda12345 · 2025-07-11T12:17:18Z

vllm/v1/executor/abstract.py

@@ -73,10 +73,16 @@ def register_failure_callback(self, callback: FailureCallback):
        pass

    def determine_available_memory(self) -> list[int]:  # in bytes
+        if self.vllm_config.model_config.is_attention_free:


Can you put the change in determine_available_memory to gpu_worker and the change inget_kv_cache_specs to gpu_model_runner? Though these changes are platform independent, this abstract class should be as simple as possible.

heheda12345 · 2025-07-11T14:52:01Z

Seems that someone want prefix caching to be enabled in "last" pooling method. Is it possible to only use the zero-group code path for other pooling methods?

vllm/vllm/config.py

Line 4594 in 8aeaa91

if pooling_type is None or pooling_type.lower() != "last":

heheda12345 · 2025-07-14T02:18:20Z

@christian-pinto #20661 is merged. Can you rebase?

christian-pinto · 2025-07-14T10:03:17Z

Seems that someone want prefix caching to be enabled in "last" pooling method. Is it possible to only use the zero-group code path for other pooling methods?

To my understanding (I am not an expert on this) you should not be able to perform prefix caching if there is no KV cache, right? Therefore I would believe that if people mark a model as attention free they would also not require any prefix caching?
Please correct me if I am wrong.

Forces 0 KV Cache groups to disable KV Cache in attention free models Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

- Changes after vllm-project#20661 merge - Fixed one pre-commit error Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

heheda12345 · 2025-07-14T10:35:24Z

@DarkLight1337 @maxdebayser See the above comment. Should we enable prefix caching for "last" pooling method? I'm also not an expert on this.

heheda12345 · 2025-07-14T14:18:34Z

vllm/v1/core/kv_cache_coordinator.py

@@ -388,7 +388,9 @@ def get_kv_cache_coordinator(
        kv_cache_config: KVCacheConfig, max_model_len: int, use_eagle: bool,
        enable_caching: bool, caching_hash_fn: Callable,
        enable_kv_cache_events: bool) -> KVCacheCoordinator:
-    if not enable_caching:
+    if not enable_caching or len(kv_cache_config.kv_cache_groups) == 0:
+        # We instantiate this coordinator also for attention free models that


Do we need this given prefix caching is disabled here for models that don’t use last pooling method?
https://github.com/maxdebayser/vllm/blob/221f013922c0c118b682d294755e69990b2c43ed/vllm/config.py#L4505

Without this check though you would not be able to disable attention for models that are not of the pooling type as prefix caching is enabled by default for all models except pooling ones.

See below:

vllm/vllm/engine/arg_utils.py

Lines 1620 to 1630 in 38efa28

def _set_default_args_v1(self, usage_context: UsageContext,

model_config: ModelConfig) -> None:

"""Set Default Arguments for V1 Engine."""

# V1 always uses chunked prefills and prefix caching

# for non-pooling tasks.

# For pooling tasks the default is False

if model_config.runner_type != "pooling":

self.enable_chunked_prefill = True

if self.enable_prefix_caching is None:

self.enable_prefix_caching = True

Perhaps the safest thing to do is to disable prefix-caching in VllmConfig.__post_init__ right away for any attention free models and then yes, we could just rely on enable_caching as you suggest.

remove this comment?

heheda12345 · 2025-07-14T14:21:11Z

vllm/v1/core/kv_cache_manager.py

@@ -89,7 +89,7 @@ def __init__(
        self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

        self.block_size: Optional[int] = None
-        if self.enable_caching:
+        if self.enable_caching and len(kv_cache_config.kv_cache_groups) > 0:


same as above. Don’t need if enable_caching is false for attention free models.

maxdebayser · 2025-07-14T14:40:11Z

@DarkLight1337 @maxdebayser See the above comment. Should we enable prefix caching for "last" pooling method? I'm also not an expert on this.

@heheda12345 , last pooling with prefix caching should only be enabled for decoder models with kv-cache.

christian-pinto · 2025-07-14T19:54:34Z

@maxdebayser I have noticed that the logic looking after enabling/disabling chunked prefill and prefix caching for pooling models is kinda duplicated?

here:
https://github.com/maxdebayser/vllm/blob/221f013922c0c118b682d294755e69990b2c43ed/vllm/config.py#L4505-L4522

and here:
https://github.com/maxdebayser/vllm/blob/221f013922c0c118b682d294755e69990b2c43ed/vllm/engine/arg_utils.py#L1527-L1548

Is this intentional in case the VllmConfig is not initialized via create_engine_config()? Keeping all of this in the post_init of VllmConfig would anyway ensure those fields are properly initialized, right?

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

christian-pinto · 2025-07-14T20:34:12Z

@heheda12345 with fb3ecfb I disable chunked prefill and prefix aching for all attention free models. This Way the changes to the KVCacheCoordinator ad KVCacheManager are further reduced.

… retention Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

heheda12345

Thanks for your patience. Sorry for didn't noticing the bug related to mamba layers.

heheda12345 · 2025-07-15T09:58:41Z

vllm/v1/worker/gpu_worker.py

@@ -209,6 +209,7 @@ def determine_available_memory(self) -> int:
            You may limit the usage of GPU memory
            by adjusting the `gpu_memory_utilization` parameter.
        """
+


remove this?

heheda12345 · 2025-07-15T10:05:17Z

vllm/v1/core/kv_cache_coordinator.py

@@ -388,7 +388,9 @@ def get_kv_cache_coordinator(
        kv_cache_config: KVCacheConfig, max_model_len: int, use_eagle: bool,
        enable_caching: bool, caching_hash_fn: Callable,
        enable_kv_cache_events: bool) -> KVCacheCoordinator:
-    if not enable_caching:
+    if not enable_caching or len(kv_cache_config.kv_cache_groups) == 0:
+        # We instantiate this coordinator also for attention free models that


remove this comment?

heheda12345 · 2025-07-15T10:06:19Z

vllm/v1/core/kv_cache_manager.py

+        self.enable_caching = (enable_caching
+                               if len(kv_cache_config.kv_cache_groups) > 0
+                               else False)


Suggested change

self.enable_caching = (enable_caching

if len(kv_cache_config.kv_cache_groups) > 0

else False)

```suggestions

if len(kv_cache_config.kv_cache_groups) == 0:

# Attention free models don't have kv cache, thus don't need prefix caching.

enable_caching = False

self.enable_caching = enable_caching

I think this structure is more clear for readers.

heheda12345 · 2025-07-15T10:16:16Z

vllm/v1/engine/core.py

+        check_available_memory = not(len(kv_cache_specs) == 1 and not kv_cache_specs[0])
+        available_gpu_memory = [0]
+        if check_available_memory:
+            available_gpu_memory = (
+                self.model_executor.determine_available_memory())


Suggested change

check_available_memory = not(len(kv_cache_specs) == 1 and not kv_cache_specs[0])

available_gpu_memory = [0]

if check_available_memory:

available_gpu_memory = (

self.model_executor.determine_available_memory())

has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)

if has_kv_cache:

available_gpu_memory = self.model_executor.determine_available_memory()

else:

# Attention free models don't need memory for kv cache

available_gpu_memory = [0] * len(kv_cache_specs)

I feel that the condition is not correct. len(kv_cache_specs) can be larger than 1 when TP / PP is enabled.

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

heheda12345

LGTM! Thank you very much.

christian-pinto · 2025-07-15T10:59:14Z

LGTM! Thank you very much.

Super. Many thanks for your help.

christian-pinto requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners July 11, 2025 10:55

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

mergify bot added the v1 label Jul 11, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

christian-pinto commented Jul 11, 2025

View reviewed changes

christian-pinto force-pushed the attention_free_models_support branch from 4b0a9f9 to e06e264 Compare July 11, 2025 11:03

heheda12345 reviewed Jul 11, 2025

View reviewed changes

christian-pinto force-pushed the attention_free_models_support branch from bb9f2db to 177b788 Compare July 14, 2025 08:31

christian-pinto added 4 commits July 14, 2025 10:04

Support for attention free models

b764c9d

Forces 0 KV Cache groups to disable KV Cache in attention free models Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

is_kv_cache_type_attention_free: return False if not attention free

5825ba4

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

some minor edits after first review round

fc86350

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Rebase to current master

97c11e6

- Changes after vllm-project#20661 merge - Fixed one pre-commit error Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

christian-pinto force-pushed the attention_free_models_support branch from 177b788 to 97c11e6 Compare July 14, 2025 10:06

Make pre-commits pass

673aeb0

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

christian-pinto mentioned this pull request Jul 14, 2025

[RFC]: Extending VLLM towards native support of non text-generating models #16052

Open

1 task

heheda12345 reviewed Jul 14, 2025

View reviewed changes

christian-pinto requested review from simon-mo and youkaichao as code owners July 14, 2025 20:16

christian-pinto requested review from mgoin, tlrmchlsmth, houseroad and hmellor as code owners July 14, 2025 20:16

Disable chunk prefill and prefix caching when model is attention free

fb3ecfb

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

christian-pinto force-pushed the attention_free_models_support branch from 0923e34 to fb3ecfb Compare July 14, 2025 20:31

christian-pinto added 3 commits July 15, 2025 09:57

reworked to allow for models like mamba to use the kv_cache for state…

8e5dbee

… retention Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

cleanup config.py

2ee7087

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

cleanup gpu_worker.py

19a7d70

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

heheda12345 reviewed Jul 15, 2025

View reviewed changes

Edits after review

b8f355e

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

heheda12345 approved these changes Jul 15, 2025

View reviewed changes

heheda12345 changed the title ~~[v1][core]Support for attention free models~~ [v1][core] Support for attention free models Jul 15, 2025

heheda12345 enabled auto-merge (squash) July 15, 2025 10:49

github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 15, 2025

heheda12345 merged commit 4ffd963 into vllm-project:main Jul 15, 2025
75 checks passed

christian-pinto mentioned this pull request Jul 16, 2025

[Core][Model] PrithviMAE Enablement on vLLM v1 engine #20577

Open

christian-pinto deleted the attention_free_models_support branch July 16, 2025 14:15

		# This returns a kv_cahce config with 0 kv cache groups and 1 block
		# to allow for the KVCache manager to handle attention fre models.

	def _set_default_args_v1(self, usage_context: UsageContext,
	model_config: ModelConfig) -> None:
	"""Set Default Arguments for V1 Engine."""

	# V1 always uses chunked prefills and prefix caching
	# for non-pooling tasks.
	# For pooling tasks the default is False
	if model_config.runner_type != "pooling":
	self.enable_chunked_prefill = True
	if self.enable_prefix_caching is None:
	self.enable_prefix_caching = True

-        self.enable_caching = (enable_caching
-                               if len(kv_cache_config.kv_cache_groups) > 0
-                               else False)
+        ```suggestions
+if len(kv_cache_config.kv_cache_groups) == 0:
+    # Attention free models don't have kv cache, thus don't need prefix caching.
+    enable_caching = False
+self.enable_caching = enable_caching

Uh oh!

[v1][core] Support for attention free models #20811

[v1][core] Support for attention free models #20811

Uh oh!

Conversation

christian-pinto commented Jul 11, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 11, 2025

Uh oh!

heheda12345 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

heheda12345 commented Jul 11, 2025

Uh oh!

heheda12345 commented Jul 14, 2025

Uh oh!

christian-pinto commented Jul 14, 2025

Uh oh!

heheda12345 commented Jul 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

christian-pinto Jul 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maxdebayser commented Jul 14, 2025

Uh oh!

christian-pinto commented Jul 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

christian-pinto commented Jul 14, 2025

Uh oh!

heheda12345 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

christian-pinto commented Jul 11, 2025 •

edited by github-actions bot

Loading

christian-pinto Jul 14, 2025 •

edited

Loading

christian-pinto commented Jul 14, 2025 •

edited

Loading